{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'pyreg'",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-1-7b887d569388>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     13\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     14\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfunctional\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mF\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 15\u001b[0;31m \u001b[0;32mfrom\u001b[0m \u001b[0mpyreg\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mutils\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mapply_affine_transform_to_map_multiNC\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mget_inverse_affine_param\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mupdate_affine_param\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     16\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mpyreg\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msimple_interface\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mSI\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     17\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mpyreg\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfileio\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mFIO\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'pyreg'"
     ],
     "output_type": "error"
    }
   ],
   "source": [
    "import sys\n",
    "import os\n",
    "os.environ[\"CUDA_DEVICE_ORDER\"] = \"PCI_BUS_ID\"\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = '1'\n",
    "os.environ[\"CUDA_CACHE_PATH\"] = \"/playpen/zyshen/summer/\"\n",
    "sys.path.append(os.path.realpath(\"..\"))\n",
    "sys.path.append(os.path.realpath(\"../mermaid\"))\n",
    "sys.path.append(os.path.realpath(\"../result_analysis\"))\n",
    "# import matplotlib as matplt\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from pyreg.utils import apply_affine_transform_to_map_multiNC, get_inverse_affine_param, update_affine_param\n",
    "import pyreg.simple_interface as SI\n",
    "import pyreg.fileio as FIO\n",
    "import pyreg.utils as pyreg_utils\n",
    "import pymesh\n",
    "from OAI_registration import OAIRegister\n",
    "\n",
    "from ipywidgets import interact, interactive, fixed, interact_manual\n",
    "import ipywidgets as widgets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "moving_img_path = '/playpen-raid/zhenlinx/Data/OAI/9010952/MR_SAG_3D_DESS/LEFT KNEE/36 MONTH/image_normalized.nii.gz'\n",
    "target_img_path = '/playpen-raid/zhenlinx/Data/OAI/9010952/MR_SAG_3D_DESS/LEFT KNEE/ENROLLMENT/image_normalized.nii.gz'\n",
    "FC_mesh_path = '/playpen-raid/zhenlinx/Data/OAI/9010952/MR_SAG_3D_DESS/LEFT KNEE/36 MONTH/Cascaded_1_AC_residual-1-s1_end2end_multi-out_UNet_bias_Nifti_rescaled_train1_patch_128_128_32_batch_4_sample_0.01-0.02_cross_entropy_lr_0.0005_scheduler_multiStep_04282018_011610/FC_thickness.ply'\n",
    "\n",
    "oai_reg = OAIRegister([('affine_map', 'reg_affine_settings.json'),\n",
    "                       ('svf_vector_momentum_map', 'reg_svf_settings.json')])\n",
    "oai_reg.set_moving_image(moving_img_path)\n",
    "oai_reg.set_target_image(target_img_path)\n",
    "spacing = oai_reg.moving_spacing_normalized"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "forward_map = np.load('../result_analysis/test_data/svf_map.npy')\n",
    "inverse_map = np.load('../result_analysis/test_data/svf_inverse_map.npy')\n",
    "forward_map_old = np.load('../result_analysis/test_data/svf_map_old.npy')\n",
    "inverse_map_old = np.load('../result_analysis/test_data/svf_inverse_map_old.npy')\n",
    "forward_map_no_init = np.load('../result_analysis/test_data/svf_map_test.npy')\n",
    "inverse_map_no_init = np.load('../result_analysis/test_data/svf_inverse_map_test.npy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "forward_map_affine = np.load('../result_analysis/test_data/affine_map.npy')\n",
    "inverse_map_affine = np.load('../result_analysis/test_data/affine_inverse_map.npy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(3, 160, 384, 384)\n",
      "[ 0.00501301  0.00261097  0.00261097]\n"
     ]
    }
   ],
   "source": [
    "map_shape = circle_map.shape\n",
    "print(map_shape)\n",
    "print(spacing)\n",
    "x=np.arange(0,map_shape[2])\n",
    "y=np.arange(0,map_shape[3])\n",
    "X, Y = np.meshgrid(x,y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "forward_map_affine200svf1 = np.load('../result_analysis/test_data/affine200svf1/svf_map.npy')\n",
    "inverse_map_affine200svf1 = np.load('../result_analysis/test_data/affine200svf1/svf_inverse_map.npy')\n",
    "forward_map_affine200svf20 = np.load('../result_analysis/test_data/affine200svf20/svf_map.npy')\n",
    "inverse_map_affine200svf20 = np.load('../result_analysis/test_data/affine200svf20/svf_inverse_map.npy')\n",
    "forward_map_affine1svf20 = np.load('../result_analysis/test_data/affine1svf20/svf_map.npy')\n",
    "inverse_map_affine1svf20 = np.load('../result_analysis/test_data/affine1svf20/svf_inverse_map.npy')\n",
    "forward_map_affine1svf1 = np.load('../result_analysis/test_data/affine1svf1/svf_map.npy')\n",
    "inverse_map_affine1svf1 = np.load('../result_analysis/test_data/affine1svf1/svf_inverse_map.npy')\n",
    "forward_map_affine0svf20 = np.load('../result_analysis/test_data/affine0svf20/svf_map.npy')\n",
    "inverse_map_affine0svf20 = np.load('../result_analysis/test_data/affine0svf20/svf_inverse_map.npy')\n",
    "\n",
    "forward_map_affine200 = np.load('../result_analysis/test_data/affine200svf20/affine_map.npy')\n",
    "inverse_map_affine200 = np.load('../result_analysis/test_data/affine200svf20/affine_inverse_map.npy')\n",
    "circle_map_affine200svf1 = pyreg_utils.compute_warped_image_multiNC(torch.from_numpy(forward_map_affine200svf1).cuda(),\n",
    "                                                      torch.from_numpy(inverse_map_affine200svf1).cuda(), \n",
    "                                                      spacing, 1,False).cpu().data.numpy().squeeze()/(spacing.reshape(3,1,1,1))\n",
    "circle_map_affine200svf20 = pyreg_utils.compute_warped_image_multiNC(torch.from_numpy(forward_map_affine200svf20).cuda(),\n",
    "                                                      torch.from_numpy(inverse_map_affine200svf20).cuda(), \n",
    "                                                      spacing, 1,False).cpu().data.numpy().squeeze()/(spacing.reshape(3,1,1,1))\n",
    "\n",
    "circle_map_affine1svf20 = pyreg_utils.compute_warped_image_multiNC(torch.from_numpy(forward_map_affine1svf20).cuda(),\n",
    "                                                      torch.from_numpy(inverse_map_affine1svf20).cuda(), \n",
    "                                                      spacing, 1,False).cpu().data.numpy().squeeze()/(spacing.reshape(3,1,1,1))\n",
    "circle_map_affine0svf20 = pyreg_utils.compute_warped_image_multiNC(torch.from_numpy(forward_map_affine0svf20).cuda(),\n",
    "                                                      torch.from_numpy(inverse_map_affine0svf20).cuda(), \n",
    "                                                      spacing, 1,False).cpu().data.numpy().squeeze()/(spacing.reshape(3,1,1,1))\n",
    "circle_map_affine1svf1 = pyreg_utils.compute_warped_image_multiNC(torch.from_numpy(forward_map_affine1svf1).cuda(),\n",
    "                                                      torch.from_numpy(inverse_map_affine1svf1).cuda(), \n",
    "                                                      spacing, 1,False).cpu().data.numpy().squeeze()/(spacing.reshape(3,1,1,1))\n",
    "circle_map_affine200 = pyreg_utils.compute_warped_image_multiNC(torch.from_numpy(forward_map_affine200).cuda(),\n",
    "                                                      torch.from_numpy(inverse_map_affine200).cuda(), \n",
    "                                                      spacing, 1,False).cpu().data.numpy().squeeze()/(spacing.reshape(3,1,1,1))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def show_circle_maps(slice_ind, circle_map_affine200svf1, circle_map_affine200svf20, circle_map_affine1svf20,\n",
    "                     circle_map_affine0svf20, circle_map_affine1svf1, circle_map_affine200):\n",
    "    figs, axs = plt.subplots(2,3, figsize=[15,10])\n",
    "\n",
    "    CSX = axs[0,0].contour(X, Y, circle_map_affine200[1,slice_ind,:,:], colors=['white'])\n",
    "    axs[0,0].imshow(circle_map_affine200[1,slice_ind,:,:])\n",
    "    axs[0,0].clabel(CSX, inline=1, fontsize=10)\n",
    "    axs[0,0].set_title('affine200\\n X coutours ')\n",
    "\n",
    "    CSY = axs[1,0].contour(X, Y, circle_map_affine200[2,slice_ind,:,:], colors=['white'])\n",
    "    axs[1,0].imshow(circle_map_affine200[2,slice_ind,:,:])\n",
    "    axs[1,0].clabel(CSY, inline=1, fontsize=10)\n",
    "    axs[1,0].set_title('Y coutours')\n",
    "\n",
    "    CSX = axs[0,1].contour(X, Y, circle_map_affine200svf1[1,slice_ind,:,:], colors=['white'])\n",
    "    axs[0,1].imshow(circle_map_affine200svf1[1,slice_ind,:,:])\n",
    "    axs[0,1].clabel(CSX, inline=1, fontsize=10)\n",
    "    axs[0,1].set_title('affine200svf1\\n X coutours')\n",
    "\n",
    "    CSY = axs[1,1].contour(X, Y, circle_map_affine200svf1[2,slice_ind,:,:], colors=['white'])\n",
    "    axs[1,1].imshow(circle_map_affine200svf1[2,slice_ind,:,:])\n",
    "    axs[1,1].clabel(CSY, inline=1, fontsize=10)\n",
    "    axs[1,1].set_title('Y coutours')\n",
    "\n",
    "    CSX = axs[0,2].contour(X, Y, circle_map_affine200svf20[1,slice_ind,:,:], colors=['white'])\n",
    "    axs[0,2].imshow(circle_map_affine200svf20[1,slice_ind,:,:])\n",
    "    axs[0,2].clabel(CSX, inline=1, fontsize=10)\n",
    "    axs[0,2].set_title('affine200svf20\\n X coutours')\n",
    "\n",
    "    CSY = axs[1,2].contour(X, Y, circle_map_affine200svf20[2,slice_ind,:,:], colors=['white'])\n",
    "    axs[1,2].imshow(circle_map_affine200svf20[2,slice_ind,:,:])\n",
    "    axs[1,2].clabel(CSY, inline=1, fontsize=10)\n",
    "    axs[1,2].set_title('Y coutours')\n",
    "\n",
    "\n",
    "    figs, axs = plt.subplots(2,3, figsize=[15,10])\n",
    "\n",
    "\n",
    "\n",
    "    CSX = axs[0,0].contour(X, Y, circle_map_affine0svf20[1,slice_ind,:,:], colors=['white'])\n",
    "    axs[0,0].imshow(circle_map_affine0svf20[1,slice_ind,:,:])\n",
    "    axs[0,0].clabel(CSX, inline=1, fontsize=10)\n",
    "    axs[0,0].set_title('affine0svf20\\n X coutours')\n",
    "\n",
    "    CSY = axs[1,0].contour(X, Y, circle_map_affine0svf20[2,slice_ind,:,:], colors=['white'])\n",
    "    axs[1,0].imshow(circle_map_affine0svf20[2,slice_ind,:,:])\n",
    "    axs[1,0].clabel(CSY, inline=1, fontsize=10)\n",
    "    axs[1,0].set_title('Y coutours')\n",
    "\n",
    "    CSX = axs[0,1].contour(X, Y, circle_map_affine1svf1[1,slice_ind,:,:], colors=['white'])\n",
    "    axs[0,1].imshow(circle_map_affine1svf1[1,slice_ind,:,:])\n",
    "    axs[0,1].clabel(CSX, inline=1, fontsize=10)\n",
    "    axs[0,1].set_title('affine1svf1 \\n X coutours')\n",
    "\n",
    "    CSY = axs[1,1].contour(X, Y, circle_map_affine1svf1[2,slice_ind,:,:], colors=['white'])\n",
    "    axs[1,1].imshow(circle_map_affine1svf1[2,slice_ind,:,:])\n",
    "    axs[1,1].clabel(CSY, inline=1, fontsize=10)\n",
    "    axs[1,1].set_title('Y coutours')\n",
    "\n",
    "    CSX = axs[0,2].contour(X, Y, circle_map_affine1svf20[1,slice_ind,:,:], colors=['white'])\n",
    "    axs[0,2].imshow(circle_map_affine1svf20[1,slice_ind,:,:])\n",
    "    axs[0,2].clabel(CSX, inline=1, fontsize=10)\n",
    "    axs[0,2].set_title('affine1svf20\\n X coutours')\n",
    "\n",
    "    CSY = axs[1,2].contour(X, Y, circle_map_affine1svf20[2,slice_ind,:,:], colors=['white'])\n",
    "    axs[1,2].imshow(circle_map_affine1svf20[2,slice_ind,:,:])\n",
    "    axs[1,2].clabel(CSY, inline=1, fontsize=10)\n",
    "    axs[1,2].set_title('Y coutours')\n",
    "\n",
    "index_slicer = widgets.IntSlider(min=0,max=159,step=1,value=80)\n",
    "out = widgets.interact(show_circle_maps, slice_ind = index_slicer, circle_map_affine200svf1=fixed(circle_map_affine200svf1), \n",
    "                                 circle_map_affine200svf20=fixed(circle_map_affine200svf20), \n",
    "                                 circle_map_affine1svf20=fixed(circle_map_affine1svf20),\n",
    "                                 circle_map_affine0svf20=fixed(circle_map_affine0svf20),\n",
    "                                 circle_map_affine1svf1=fixed(circle_map_affine1svf1),\n",
    "                                 circle_map_affine200=fixed(circle_map_affine200))\n",
    "display(out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "figs, axs = plt.subplots(2,2, figsize=[10,10])\n",
    "CSX = axs[0,0].contour(X, Y, forward_map_debug.squeeze()[1,slice_ind,:,:]/spacing[1])\n",
    "axs[0,0].clabel(CSX, inline=1, fontsize=10)\n",
    "axs[0,0].set_title('X coutours')\n",
    "\n",
    "CSY = axs[1,0].contour(X, Y, forward_map_debug.squeeze()[2,slice_ind,:,:]/spacing[2])\n",
    "axs[1,0].clabel(CSY, inline=1, fontsize=10)\n",
    "axs[1,0].set_title('Y coutours')\n",
    "\n",
    "\n",
    "CSX = axs[0,1].contour(X, Y, inverse_map_debug.squeeze()[1,slice_ind,:,:]/spacing[1])\n",
    "axs[0,1].clabel(CSX, inline=1, fontsize=10)\n",
    "axs[0,1].set_title('X coutours')\n",
    "\n",
    "CSY = axs[1,1].contour(X, Y, inverse_map_debug.squeeze()[2,slice_ind,:,:]/spacing[2])\n",
    "axs[1,1].clabel(CSY, inline=1, fontsize=10)\n",
    "axs[1,1].set_title('Y coutours')\n",
    "\n",
    "\n",
    "figs, axs = plt.subplots(2,2, figsize=[10,10])\n",
    "CSX = axs[0,0].contour(X, Y, forward_map_no_init.squeeze()[1,slice_ind,:,:]/spacing[1])\n",
    "axs[0,0].clabel(CSX, inline=1, fontsize=10)\n",
    "axs[0,0].set_title('X coutours')\n",
    "\n",
    "CSY = axs[1,0].contour(X, Y, forward_map_no_init.squeeze()[2,slice_ind,:,:]/spacing[2])\n",
    "axs[1,0].clabel(CSY, inline=1, fontsize=10)\n",
    "axs[1,0].set_title('Y coutours')\n",
    "\n",
    "\n",
    "CSX = axs[0,1].contour(X, Y, inverse_map_no_init.squeeze()[1,slice_ind,:,:]/spacing[1])\n",
    "axs[0,1].clabel(CSX, inline=1, fontsize=10)\n",
    "axs[0,1].set_title('X coutours')\n",
    "\n",
    "CSY = axs[1,1].contour(X, Y, inverse_map_no_init.squeeze()[2,slice_ind,:,:]/spacing[2])\n",
    "axs[1,1].clabel(CSY, inline=1, fontsize=10)\n",
    "axs[1,1].set_title('Y coutours')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
